#include "../util.h"
#include "../translation_model/tm.h"

#include <algorithm>
#include <cassert>
#include <cmath>
#include <fstream>
#include <functional>
#include <iostream>
#include <string>

using std::string;
using std::vector;
using std::ifstream;
using std::pair;

typedef pair<size_t, size_t> TInterval; // два числа - индексы токенов в предложении, определяющие фразу
typedef pair<TInterval, TInterval> TIntervalAlignment; // два сопоставленных друг другу интервала

// Регулирует размер штрафа, назначаемого за оценку слову по более общему контексту. Чем больше, тем сильнее штраф.
const double NGRAM_LENGTH_PENALTY_COEFF = 0.9;
// Регулирует соотношение между оценками, полученным по моделям языка и перевода. 
// Численно равен степени alpha, в которую возводится вероятность P(r) в формуле argmax P(e|r)*P(r)^alpha
const double LANG_MODEL_COEFF = 0.5;
// Чем меньше значение, тем в переводе было меньше число фраз.
// Несколько нивелирует эффект того, что модель языка сильно штрафует хорошие длинные переводы,
// где каждая фраза передается большим числом слов.
//default 1.0
const double PHRASE_COUNT_PENALTY = -0.9;

const double N_GRAM_COEFF[4] = {0.0, 0.1, 0.4, 1.0};
const bool GOOD_TURING_SMOOTHING = true;

//const double NEED_SMOOTH = 0.01;
//const int NEED_SMOOTH_CNT = 20;

// Информация об одном из переводов фразы
struct TCandidateInfo {
    // токенизированное предложение
    vector<string> Tokens;
    // проставленное соответствие между интервалами
    vector<TIntervalAlignment> Alignment;
};

void PrintUsage() {
    std::cerr << "Usage: baseline config\n";
}

void Log(const string& msg) {
    std::cerr << msg << '\n';
}

// прочитать из файла ngramFilename n-граммы в частотный словарь ngrams
void ReadNgramFile(const string& ngramFilename, TFreqDict& ngrams) {
    ifstream ngramIfstream(ngramFilename.c_str());
    string line;
    while (getline(ngramIfstream, line)) {
        vector<string> tokens = SplitString(line, '\t');
        assert(tokens.size() == 2);
        string ngram = tokens[0];
        size_t count = FromString<size_t>(tokens[1]);
		ngrams.dict[ngram] = count;
		ngrams.r[count]++;
    }
}

// заполнить информацию о соответствии двух интервалов по строке вида 3-5=2-3
TInterval ParseInterval(const string& s) {
    if (s.find('-') == string::npos) {
        size_t value = FromString<size_t>(s);
        return TInterval(value, value);
    }
    vector<string> tokens = SplitString(s, '-');
    assert(tokens.size() == 2);
    return TInterval(FromString<size_t>(tokens[0]), FromString<size_t>(tokens[1]));
}

// заполнить информацию о соответствии фраз исходного предложения и перевода
void ParseAlignment(const string& s, vector<TIntervalAlignment>& alignments) {
    vector<string> alignmentPairs = SplitString(s);
    size_t alignmentCount = alignmentPairs.size();
    alignments.resize(alignmentCount);
    for (size_t i = 0; i < alignmentCount; ++i) {
        vector<string> pair = SplitString(alignmentPairs[i], '=');
        assert(pair.size() == 2);
        alignments[i] = TIntervalAlignment(ParseInterval(pair[0]), ParseInterval(pair[1]));
    }
}

// подсчитать оценку вероятности появления предложения words в языке, где язык задан своими n-граммами ngrams
double CalculateLangModelScore(const vector<string>& words, const vector<TFreqDict>& ngrams) {
    size_t ngramMaxLength = ngrams.size();
    // Добавим к словам предложения токены, отмечающие начало и конец предложения - <s> и </s>, соответственно.
    vector<string> vs;
    vs.push_back("<s>");
    vs.insert(vs.end(), words.begin(), words.end());
    vs.push_back("</s>");
    double result = 0.0;
    // Пройти по всем словам предложения (исключая токены начала и конца предложения - <s> и </s>)
    // и подсчитать вероятность того, что такое предложение встретится в языке как произведение вероятностей встретить
    // каждое слово. Для того, чтобы получить сумму чисел, берем -ln(произведения). Меньшее значение соответсвует большей
    // вероятности предложения, само значение неотрицательно.
    for (size_t currWordPos = 1; currWordPos + 1 < vs.size(); ++currWordPos) {
        // ngramWithWord - n-грамма, заканчивающаяся на текущем слове
        // ngramWithoutWord - (n-1)-грамма, полученная из ngramWithWord удалением последнего слова
        // например: предложение <s> a b c </s>, текущее слово b (currWordPos = 2), длина n-грамм = 3.
        // тогда ngramWithWord = "<s> a b", а ngramWithoutWord = "<s> a"
        // Вероятность встретить в языке в данном контексте токен b оценивается как freq("<s> a b")/freq("<s> a")

        size_t ngramStartLength = std::min(currWordPos + 1, ngramMaxLength);
        double wordLangScore = 0.0;
		bool flag = false;
        for (size_t ngramLength = ngramStartLength; ngramLength >= 1; --ngramLength) {
            size_t ngramStartPos = currWordPos - (ngramLength - 1);

            // JoinString для итераторов склеивает слова, не включая правую границу
            string ngramWithoutWord = JoinString(vs.begin() + ngramStartPos, vs.begin() + currWordPos);
            string ngramWithWord = JoinString(vs.begin() + ngramStartPos, vs.begin() + currWordPos + 1);

            // Индексы массива ngrams начинаются с нуля, так что ngrams[0] содержит униграммы, ngrams[1] - биграммы и т.д.
            TFreqDictPart::const_iterator it = ngrams[ngramLength - 1].dict.find(ngramWithWord);
            if (it != ngrams[ngramLength - 1].dict.end()) {
                size_t freqNgramWithWord = it->second;
                // Если длина n-граммы равна 1, то контекста для нее нет.
                // Тогда за частоту контекста принимается сумма частот всех униграмм.
                TFreqDictPart::const_iterator ngramWithoutWordIterator;
                if (ngramLength == 1) {
                    ngramWithoutWordIterator = ngrams[0].dict.find(" ");
                    assert(ngramWithoutWordIterator != ngrams[0].dict.end());
                } else {
                    ngramWithoutWordIterator = ngrams[ngramLength - 2].dict.find(ngramWithoutWord);
                    assert(ngramWithoutWordIterator != ngrams[ngramLength - 2].dict.end());
                }
                size_t freqNgramWithoutWord = ngramWithoutWordIterator->second;

                // если мы получили оценку по короткому контексту (малая длина n-граммы), то штраф
                double ngramLengthPenalty = NGRAM_LENGTH_PENALTY_COEFF * (ngramMaxLength - ngramLength);

				TFreqCounter::const_iterator freq_it1 = ngrams[ngramLength-1].r.find(freqNgramWithWord);
				TFreqCounter::const_iterator freq_it2 = ngrams[ngramLength-1].r.find(freqNgramWithWord+1);

				double r_0 = 1.0*(freq_it1->second);
				double r_1 = (freq_it2==ngrams[ngramLength-1].r.end())?1.0:1.0*(freq_it2->second);

				double p_0 = 1.0;
				double p_1 = 1.0;
				if (ngramLength!=1)
				{
					freq_it1 = ngrams[ngramLength-2].r.find(freqNgramWithoutWord);
					freq_it2 = ngrams[ngramLength-2].r.find(freqNgramWithoutWord+1);
					p_0 = 1.0*(freq_it1->second);
					p_1 = (freq_it2==ngrams[ngramLength-2].r.end())?1.0:1.0*(freq_it2->second);
				} 				

				double a = freqNgramWithWord, b = freqNgramWithoutWord;
				
				if (GOOD_TURING_SMOOTHING)
				{
					a = (freqNgramWithWord+1.0)*r_1/r_0;
					b = (freqNgramWithoutWord+1.0)*p_1/p_0;
				} 
				wordLangScore += N_GRAM_COEFF[ngramLength] * a/b;
				flag = true;
            }
        }
		if (flag)
			result+=-log(wordLangScore);
    }	
//	if (vs.size()>2)
//		result /= (vs.size()-2);
    return result; // меньше - лучше
}

// подсчитать оценку того, что перевод translation является корректным для предложения на исходном языке src
// при том, что известна модель перевода - сопоставленные фразы phraseMatching
double CalculateTranslationModelScore(const vector<string>& src, const TCandidateInfo& translation,
                                      const TPhraseMatching& phraseMatching) {
    double result = 0.0;
    // по всем выровненным фразам
    for (size_t i = 0; i < translation.Alignment.size(); ++i) {
        // фраза на исходном языке...
        string enPhrase = JoinString(src.begin() + translation.Alignment[i].first.first,
                                 src.begin() + translation.Alignment[i].first.second + 1);
        // и соответствующая ей фраза на языке перевода
        string ruPhrase = JoinString(translation.Tokens.begin() + translation.Alignment[i].second.first,
                               translation.Tokens.begin() + translation.Alignment[i].second.second + 1);
        
        map<string, TFreqDict>::const_iterator ruIterator = phraseMatching.Translations.find(ruPhrase);
        if (ruIterator != phraseMatching.Translations.end()) {
            TFreqDictPart::const_iterator enIterator = ruIterator->second.dict.find(enPhrase);
            if (enIterator != ruIterator->second.dict.end()) {
                // эта пара фраз есть в собранной модели перевода
                size_t enFreq = phraseMatching.EnPhraseFreq.dict.find(enPhrase)->second;
                size_t ruFreq = phraseMatching.RuPhraseFreq.dict.find(ruPhrase)->second;
                size_t freq = enIterator->second;
                result += -0.5 * (log(double(freq) / enFreq) + log(double(freq) / ruFreq)) + PHRASE_COUNT_PENALTY;
            }
        }
    }
    return result; // меньше - лучше
}

// класс, собирающий статистику min/avg/max для числовой последовательности
// может оказаться полезным при подборе коэффициентов для моделей
class TStatInfo {
private:
    double avgValue;
    double minValue;
    double maxValue;
    double valueSum;
    size_t valueCount;

public:
    TStatInfo() {
        valueCount = 0;
    }
    void Update(double value) {
        if (valueCount == 0) {
            minValue = maxValue = value;
            valueSum = 0;
        } else {
            minValue = std::min(minValue, value);
            maxValue = std::max(maxValue, value);
        }
        valueSum += value;
        ++valueCount;
    }
    double GetAverage() const { return valueSum / valueCount; }
    double GetMin() const { return minValue; }
    double GetMax() const { return maxValue; }
};

std::ostream& operator<< (std::ostream& os, const TStatInfo& stats) {
    os << "min = " << stats.GetMin() << ", avg = " << stats.GetAverage() << ", max = " << stats.GetMax();
    return os;
}

// для строки src на исходном языке выбрать лучший перевод из кандидатов в массиве translations
// 
size_t ChooseBestTranslation(const string& src, const vector<TCandidateInfo> translations,
                             const TPhraseMatching& phraseMatching, const vector<TFreqDict>& ngrams,
                             TStatInfo& langStats, TStatInfo& translationStats) {
    double bestScore;
    size_t bestIndex;
    vector<string> srcTokens = SplitString(src);
    
    for (size_t i = 0; i < translations.size(); ++i) {
        double langScore = CalculateLangModelScore(translations[i].Tokens, ngrams);
        double translationScore = CalculateTranslationModelScore(srcTokens, translations[i], phraseMatching);
        double score = LANG_MODEL_COEFF * langScore + translationScore;
        if ((i == 0) || (score < bestScore)) {
            bestIndex = i;
            bestScore = score;
        }

        langStats.Update(LANG_MODEL_COEFF * langScore);
        translationStats.Update(translationScore);
    }
    return bestIndex;
}

void ReadConfig(const string& filename, vector<TFreqDict>& ngrams, TPhraseMatching& phraseMatching,
                vector<string>& srcSentences, vector< vector<TCandidateInfo> >& targets) {
    ifstream ifs(filename.c_str());

    Log("Reading ngrams");
    size_t ngramsMaxLevel;
    ifs >> ngramsMaxLevel;
    if (ngramsMaxLevel < 1 || ngramsMaxLevel > 5) {
        throw TException("ngramsMaxLevel should be between 1 and 5, inclusively");
    }
    ngrams.resize(ngramsMaxLevel);
    vector<string> ngramFilenames(ngramsMaxLevel);
    for (size_t i = 0; i < ngramsMaxLevel; ++i) {
        string ngramFilename;
        ifs >> ngramFilename;
        ReadNgramFile(ngramFilename, ngrams[i]);
    }

    Log("Reading translation model");
    string translationModelFilename;
    ifs >> translationModelFilename;
    ifstream translationModelIfstream(translationModelFilename.c_str());
    translationModelIfstream >> phraseMatching;

    Log("Reading input");
    string inputPhrasesFilename, translationCandidatesFilename;
    ifs >> inputPhrasesFilename >> translationCandidatesFilename;
    ifstream sourceIfstream(inputPhrasesFilename.c_str());
    string sentence;
    while (getline(sourceIfstream, sentence)) {
        srcSentences.push_back(sentence);
    }
    size_t sentenceCount = srcSentences.size();

    Log("Reading translation candidates");
    ifstream targetIfstream(translationCandidatesFilename.c_str());
    targets.resize(sentenceCount);
    string targetLine;
    while (getline(targetIfstream, targetLine)) {
        vector<string> tokens = SplitString(targetLine, " ||| ");
        assert(tokens.size() == 3);
        size_t sentenceNum = FromString<size_t>(tokens[0]);
        string candidatePhrase = tokens[1];
        RTrim(candidatePhrase);
        string alignment = tokens[2];
        TCandidateInfo candidate;
        candidate.Tokens = SplitString(candidatePhrase);
        ParseAlignment(alignment, candidate.Alignment);
        targets[sentenceNum].push_back(candidate);
    }
}

int main(int argc, char **argv) {
    if (argc != 2) {
        PrintUsage();
        return 1;
    }

    string configFilename(argv[1]);

    vector<TFreqDict> ngrams;
    TPhraseMatching phraseMatching;
    vector<string> srcSentences;
    vector< vector<TCandidateInfo> > targets;

    ReadConfig(configFilename, ngrams, phraseMatching, srcSentences, targets);

    Log("Main translation loop");
    TStatInfo langStats;
    TStatInfo translationStats;
    
    size_t sentenceCount = srcSentences.size();
    for (size_t i = 0; i < sentenceCount; ++i) {
        size_t bestTranslationIndex = ChooseBestTranslation(srcSentences[i], targets[i], phraseMatching, ngrams,
                                                            langStats, translationStats);
        // size_t bestTranslationIndex = rand() % targets[i].size();
        // size_t bestTranslationIndex = 0;
        std::cout << bestTranslationIndex << '\n';
    }

    Log("Done");
    std::cerr << "language model: " << langStats << '\n';
    std::cerr << "translation model: " << translationStats << '\n';

    return 0;
}
